#!/usr/bin/env python

import sys
import os
import time
from ftplib import FTP, error_temp, error_perm

try:
    import paramiko
except ImportError:
    paramiko = None

class IrisFTPException(Exception):
    """Super class for IrisFTPExceptions"""
    def __init__(self, msg):
        self.msg = msg
    def __str__(self):
        return self.msg

class ConnectionException(IrisFTPException):
    """Connection error."""
    def __init__(self, msg):
        self.msg = msg

class NoConnectionException(IrisFTPException):
    """No connection present."""
    def __init__(self, msg):
        self.msg = msg

class FileExistsException(IrisFTPException):
    """File allready existing."""
    def __init__(self, msg):
        self.msg = msg

class NoFileException(IrisFTPException):
    """No file found."""
    def __init__(self, msg):
        self.msg = msg

class IrisFTP(FTP):
    """Extension of the default FTP class."""
    host = None     # ftp host adress
    port = 21       # ftp port
    username = None # ftp user
    password = None # ftp password
    verbose = False # verbose mode toggle
    
    dirlist = []    # file list of current dir, element style: (type, filename)
                    # type is 'd' for directories and '-' for files

    welcome = None  # welcome message of ftp server
    connected = False   # connected flag

    # dictionary index from ls listing
    MODE,COUNT,OWNER,GROUP,SIZE,MONTH,DATE,TIME,NAME = range(0,9) 


    def __init__(self, host=None, username=None, password=None, verbose=False):
        """IrisFTP constructor.

        Arguments:
        host -- host adress
        user -- username, password tuple
        verbose -- enter verbose mode when True
        """

        FTP.__init__( self )

        self.host = host
        self.username = username
        self.password = password
        self.verbose = verbose

        if self.verbose: print "Creating IrisFTP object"

    def cwd(self, dir):
        """Change working directory."""
        if dir and not dir == '.':
            try: 
                FTP.cwd(self, dir)
            except error_perm:
                raise NoFileException, "No directory %s on server" % dir
        self.__setDirlist()

    def connect(self):
        """Connect to FTP server using current settings.

        Returns:
        Welcome message of the server
        """
        if self.connected: return

        if self.verbose: 
            print "Connecting to %s as %s" % (self.host, self.username)

        try:
            FTP.connect( self, self.host, self.port )
            self.login( self.username, self.password )
        except Exception, msg:
            raise ConnectionException, \
              "Connection denied to %s@%s" % (self.username, self.host)

        self.connected = True
        self.welcome = self.getwelcome()

        self.__setDirlist()

        return self.welcome

    def disconnect(self):
        """Disconnect from the FTP server."""
        if self.connected:
            self.quit()
            self.connected = False
            if self.verbose: print "Disconnected from %s" % self.host

    def dir(self, path=None):
        """Return directory listing.

        Arguments:
        path -- path to list, default=current dir
        """
        currentdir = self.pwd()
        self.cwd(path)
        result = self.dirlist
        self.cwd(currentdir)

        return result

    def exists(self, path):
        """Check if path exists.
        
        Arguments:
        path -- path to check existence for
        """
        if not self.connected: raise NoConnectionException, "No ftp connection present"

        (remotedir, remote) = os.path.split(path)

        currentdir = self.pwd()

        res = False

        if not remotedir: remotedir = "."

        try:
            self.cwd(remotedir)
            if remote in self.dirlist:
                res = True
        finally:
            self.cwd(currentdir)
            
        return res

    def get(self, remotefile, localfile=None, force=False):
        """Download file from remote server.

        Arguments:
        remotefile -- path to remotefile
        localfile -- local storage, default is same basename as remotefile
        force -- force overwrite of localfile, default is False
        """

        if not self.connected: raise NoConnectionException, "No ftp connection present"

        (remotedir, remote) = os.path.split(remotefile)

        if not localfile: localfile = remote

        if os.path.isdir( localfile ): localfile = os.path.join( localfile, remote )

        if os.path.exists(localfile):
            if not force:
                if self.verbose: print "%s allready exsists" % localfile
                raise FileExistsException, "Local file, %s, allready exsists" % localfile
            else:
                if self.verbose: print "Deleting %s" % localfile

                os.remove(localfile)

        currentdir = self.pwd()

        local = None
        try:
            if self.verbose: print "downloading %s to %s..." % (remotefile, localfile)

            local = open(localfile, "wb")

            self.cwd(remotedir)
            self.retrbinary('RETR ' + remote, local.write, 1024)

            if self.verbose: print "\tdownload complete"
        finally:
            if local: local.close()
            self.cwd(currentdir)

    def put(self, localfile, remotefile=None, force=False):
        """Upload file to remote server.

        Arguments:
        localfile -- path to localfile
        remotefile -- remote storage, default is same basename as localfile
        force -- force overwrite of remotefile, default is False
        """

        if not self.connected: raise NoConnectionException, "No ftp connection present"

        if not remotefile: remotefile = os.path.basename( localfile )


        try:
            self.dir(remotefile)
            if force:
                if self.verbose: print "Deleting %s" % remotefile
                self.delete(self, remotefile)
            else:
                if self.verbose: print "%s allready exsists" % localfile
                raise FileExistsException, "Local file, %s, allready exsists" % localfile
        except NoFileException, e:
            pass


        remotedir = os.path.split(remotefile)[0]

        currentdir = self.pwd()

        local = None
        try:
            if self.verbose: print "Uploading %s to %s" % (localfile, remotefile)

            local = open(localfile, "rb")
            self.cwd(remotedir)
            self.storbinary('STOR ' + remotefile, local, 1024)

            if self.verbose: print "Upload complete"

        finally:
            if local: local.close()
            self.cwd(currentdir)

    def getdir(self, remotedir, localdir=None, force=False, recursive=False):
        """Download remote directory, recursivly if specified.
        
        Arguments:
        remotedir -- remotedir to download
        localdir -- local path to store files
        force -- overwrite local files(True/False)
        recursive -- recursive download(True/False)
        """

        if not self.connected: raise NoConnectionException, "No ftp connection present"

        if remotedir[-1] == "/": remotedir = remotedir[0:-1]

        path = "/".join(remotedir.split("/")[0:-1])
        if path == "/": path = "./"
        remotedir = remotedir.split("/")[-1]

        if localdir == None: localdir = remotedir

        if os.path.exists(localdir):
            if not force:
                if self.verbose: print "%s allready exsists" % localdir
                raise FileExistsException, "Local file, %s, allready exsists" % localdir
        else:
            try:
                os.mkdir(localdir)
            except Exception, e:
                print e
                return 1


        if self.verbose: print "downloading %s/%s to %s" % (path, remotedir, localdir)

        currentdir = self.pwd()

        try:
            self.cwd(os.path.join(path,remotedir))
            for file in self.dirlist:
                l = []
                FTP.dir(self,file, lambda x: l.append(x))
                if len(l) > 1:
                    if recursive:
                        self.getdir(file, os.path.join(localdir,file), force, recursive)
                else:
                    try:
                        self.get(file, os.path.join(localdir,file), force)
                    except NoFileException, e:
                        print e
                    except error_perm, e:
                        print e
                    except FileExistsException, e:
                        print e
        finally:
            self.cwd(currentdir)

    def putdir(self, localdir, remotedir=None, force=False, recursive=False):
        """Upload local directory, recursivly if specified.
        
        Arguments:
        localdir -- local directory to upload
        remotedir -- remote location to store files
        force -- overwrite local files(True/False)
        recursive -- recursive upload(True/False)
        """

        if not self.connected: raise NoConnectionException, "No ftp connection present"

        if localdir[-1] == "/": localdir = localdir[0:-1]

        if remotedir == None: remotedir = localdir
        
        path = "/".join(remotedir.split("/")[0:-1])
        if path == "/": path = "./"
        elif path == "": path = "."
        remotedir = remotedir.split("/")[-1]

        if self.verbose: print "uploading %s/%s to %s" % (path, localdir, remotedir)

        currentdir = self.pwd()
        try:
            self.cwd(path)

            if remotedir in self.dirlist:
                if not force:
                    raise FileExistsException, "remotedir %s allready exsists" % remotedir
                else:
                    self.cwd(remotedir)
            else:
                try:
                    self.mkd(remotedir)
                    self.cwd(remotedir)
                except Exception, e:
                    print e
                    return 1

            for file in os.listdir(localdir):
                path = os.path.join(localdir,file)
                if os.path.isfile(path):
                    if self.verbose: print "uploading %s" % path
                    self.put(path, file, force)
                elif recursive and os.path.isdir(path):    
                    if self.verbose: print "uploading dir %s" % path
                    self.putdir(path, file, force, recursive)
        finally:
            self.cwd(currentdir)

    def rmdir(self, remotedir):
        """Recursively delte remote directory

        Arguments:
        remotedir -- remote directory to remove
        """

        if not self.connected: raise NoConnectionException, "No ftp connection present"

        if self.verbose: print "deleting dir %s" % remotedir

        if remotedir[-1] == "/": remotedir = remotedir[0:-1]

        path = "/".join(remotedir.split("/")[0:-1])
        if path == "/": path = "./"
        remotedir = remotedir.split("/")[-1]

        if remotedir == "": (remotedir,path) = (path,remotedir)

        currentdir = self.pwd()

        try:
            self.cwd(os.path.join(path,remotedir))
            for file in self.dirlist:
                l = []
                FTP.dir(self,file, lambda x: l.append(x))
                if len(l) > 1:
                    self.rmdir(file)
                else:
                    if self.verbose: print "deleting file %s" % file
                    self.delete(file) 
            self.cwd(os.path.join(self.pwd(), ".."))
            FTP.rmd(self, remotedir)
        finally:
            self.cwd(currentdir)

        self.__setDirlist()

    def delete(self, remotefile):
        """Delete remotefile.

        Arguments:
        remotefile -- file on the server to remove
        """

        (remotedir, file) = os.path.split(remotefile)
        if remotedir == "": remotedir = "."
        
        currentdir = self.pwd()

        self.cwd(remotedir)

        if file in self.dirlist:
            if self.verbose: print "Deleting %s" % remotefile
            FTP.delete(self, file)
        else:
            self.cwd(currentdir)
            raise NoFileException, "No file %s on server" % remotefile

        self.cwd(currentdir)

    def size(self, remotefile):
        """Get filesize of remotefile.

        Arguments:
        remotefile -- file on the server to check size of
        """
        if not self.connected: raise NoConnectionException, "No ftp connection present"

        (remotedir, remote) = os.path.split(remotefile)

        currentdir = self.pwd()

        try:
            self.cwd(remotedir)
            if remote in self.dirlist:
                l = []
                FTP.dir(self, remote, lambda x: l.append(x))
                l = l[0]
                if self.verbose: print "found file %s, size: %s" % (remote,l[self.SIZE])
                return l[self.SIZE]
        finally:
            self.cwd(currentdir)

        if self.verbose: print "File not found: %s" % remotefile

        raise NoFileException, "File not found, %s" % remotefile

    def __setDirlist(self):
        self.dirlist = [] # temporary variable to store directory listing
        FTP.dir(self, ".", \
            lambda x: self.dirlist.append(x.split()[8]))

    def time(self, remotefile):
        """Get creation time of remotefile in seconds since Epoch.
        
        Arguments:
        remotefile -- remotefile to get creation time of
        """

        monthdict = {'Jan': 1,'Feb': 2,'Mar': 3,'Apr': 4,'May': 5,'Jun': 6,'Jul': 7,'Aug': 8,'Sep': 9,'Oct': 10,'Nov': 11,'Dec':12} 

        if not self.connected: raise NoConnectionException, "No ftp connection present"

        (remotedir, remote) = os.path.split(remotefile)

        currentdir = self.pwd()

        try:
            self.cwd(remotedir)
            for file in self.dirlist:
                if remote in file:
                    if self.verbose: print "found file %s" % (remotefile)
                    remotefile = file
        finally:
            self.cwd(currentdir)

        if isinstance(remotefile, str):
            raise NoFileException, "File not found, %s" % remotefile

        month = monthdict[remotefile[self.MONTH]]
        day = int(remotefile[self.DATE])
        if len( remotefile[self.TIME]) == 5:
            year = time.localtime().tm_year
            hour = int(remotefile[self.TIME][0:2])
            min = int(remotefile[self.TIME][3:5])
        elif len( remotefile[self.TIME]) == 4:
            year = int(remotefile[self.TIME])
            hour,min = 0,0

        return time.mktime((year,month,day,hour,min,0,0,0,-1))


if paramiko:
    # class IrisSTP depens on paramiko, so only create if paramiko is loaded
    class IrisSFTP():
        """SFTP class building on the paramiko SFTPClient class"""
        host = None     # ftp host adress
        port = 22       # ftp port
        username = None # ftp user
        password = None # ftp password
        verbose = None  # verbose mode toggle
        
        dirlist = []    # file list of current dir, element style: (type, filename)
                        # type is 'd' for directories and '-' for files

        welcome = None  # welcome message of ftp server
        connected = False   # connected flag
        transport = None # transpor channel to use

        # dictionary index from ls listing
        MODE,COUNT,OWNER,GROUP,SIZE,MONTH,DATE,TIME,NAME = range(0,9) 


        def __init__(self, host=None, username=None, password=None, verbose=False):
            """IrisSFTP constructor.

            Arguments:
            host -- host adress
            user -- username, password tuple
            verbose -- enter verbose mode when True
            """

            self.host = host
            self.username = username
            self.password = password
            self.verbose = verbose

            if self.verbose: print "Creating IrisSFTP object"

        def cwd(self, dir):
            """Change working directory."""
            if dir and not dir == '.':
                try: 
                    self.sftp.chdir(dir)
                except IOError:
                    raise NoFileException, "No directory %s on server" % dir
            self.__setDirlist()

        def pwd(self):
            """Print working dirctory."""
        
            return self.sftp.getcwd() or self.sftp.normalize(".")

        def connect(self):
            """Connect to FTP server using current settings.
            
            Returns:
            Welcome message of the server
            """ 
            
            if self.connected: return

            if self.verbose: 
                print "Connecting to %s as %s" % (self.host, self.username)

            try:
                paramiko.util.log_to_file('.sftp.log')
                hostkeytype = None
                hostkey = None

                try:
                    host_keys = paramiko.util.load_host_keys(os.path.expanduser('~/.ssh/known_hosts'))
                except IOError, e:
                    try:
                        host_keys = paramiko.util.load_host_keys(os.path.expanduser('~/ssh/known_hosts'))
                    except IOError, e:
                        if self.verbose: print "Unable to load host keys file"
                        host_keys = {}
                if host_keys.has_key(self.host):
                    hostkeytype = host_keys[self.host].keys()[0]
                    hostkey = host_keys[self.host][hostkeytype]
                    if self.verbose: print "Using host key of type %s" % hostkeytype

                self.transport = paramiko.Transport((self.host, self.port))
                self.transport.connect(username=self.username, password=self.password, hostkey=hostkey)
                self.sftp = paramiko.SFTPClient.from_transport(self.transport)
            except Exception, msg:
                raise ConnectionException, \
                  "Connection denied to %s@%s" % (self.username, self.host)

            self.connected = True

            self.__setDirlist()

        def disconnect(self):
            """Disconnect from the FTP server."""
            if self.connected:
                self.transport.close()
                self.connected = False
                if self.verbose: print "Disconnected from %s" % self.host

        def dir(self, path=None):
            """Return directory listing.

            Arguments:
            path -- path to list, default=current dir
            """
            if not path: return self.dirlist

            currentdir = self.pwd()

            self.cwd(path)
            result = self.dirlist
            self.cwd(currentdir)

            return result

        def exists(self, path):
            """Check if path exists.
            
            Arguments:
            path -- path to check existence for
            """
            if not self.connected: raise NoConnectionException, "No ftp connection present"

            (remotedir, remote) = os.path.split(path)

            currentdir = self.pwd()

            res = False

            if not remotedir: remotedir = "."

            try:
                self.cwd(remotedir)
                if remote in self.dirlist:
                    res = True
            finally:
                self.cwd(currentdir)
                
            return res

        def get(self, remotefile, localfile=None, force=False):
            """Download file from remote server.

            Arguments:
            remotefile -- path to remotefile
            localfile -- local storage, default is same basename as remotefile
            force -- force overwrite of localfile, default is False
            """
            
            if not self.connected: raise NoConnectionException, "No ftp connection present"

            (remotedir, remote) = os.path.split(remotefile)

            if not localfile: localfile = remote

            if os.path.isdir( localfile ): localfile = os.path.join( localfile, remote )

            if os.path.exists(localfile):
                if not force:
                    if self.verbose: print "%s allready exsists" % localfile
                    raise FileExistsException, "Local file, %s, allready exsists" % localfile
                else:
                    if self.verbose: print "Deleting %s" % localfile

                    os.remove(localfile)

            try:
                if self.verbose: print "downloading %s to %s..." % (remotefile, localfile)

                self.sftp.get(remotefile,localfile)

                if self.verbose: print "\tdownload complete"
            except IOError:
                raise NoFileException, "Couldn't find remote file %s" % remotefile

        def put(self, localfile, remotefile=None, force=False):
            """Upload file to remote server.

            Arguments:
            localfile -- path to localfile
            remotefile -- remote storage, default is same basename as localfile
            force -- force overwrite of remotefile, default is False
            """
            
            if not self.connected: raise NoConnectionException, "No ftp connection present"

            if not remotefile: remotefile = os.path.basename( localfile )

            if self.verbose: print "Uploading %s to %s" % (localfile, remotefile)

            if self.exists(remotefile):
                if force:
                    if self.verbose: print "Deleting %s" % remotefile
                    self.delete(remotefile)
                else:
                    if self.verbose: print "%s allready exsists" % localfile
                    raise FileExistsException, "Remote file, %s, allready exsists" % localfile

            try:
                self.sftp.put(localfile, remotefile)   
            except IOError:
                raise NoFileException, "Local file, %s, does not exist" % localfile

            if self.verbose: print "Upload complete"

        def getdir(self, remotedir, localdir=None, force=False, recursive=False):
            """Download remote directory, recursivly if specified.
            
            Arguments:
            remotedir -- remotedir to download
            localdir -- local path to store files
            force -- overwrite local files(True/False)
            recursive -- recursive download(True/False)
            """

            if not self.connected: raise NoConnectionException, "No ftp connection present"

            if remotedir[-1] == "/": remotedir = remotedir[0:-1]

            path = "/".join(remotedir.split("/")[0:-1])
            if path == "/": path = "./"
            remotedir = remotedir.split("/")[-1]

            if localdir == None: localdir = remotedir

            if os.path.exists(localdir):
                if not force:
                    if self.verbose: print "%s allready exsists" % localdir
                    raise FileExistsException, "Local file, %s, allready exsists" % localdir
            else:
                try:
                    os.mkdir(localdir)
                except Exception, e:
                    print e
                    return 1


            if self.verbose: print "downloading %s/%s to %s" % (path, remotedir, localdir)

            currentdir = self.pwd()

            try:
                self.cwd(os.path.join(path,remotedir))
                for file in self.dirlist:
                    if str(self.sftp.stat(file))[0] == 'd':
                        if recursive:
                            self.getdir(file, os.path.join(localdir,file), force, recursive)
                    else:
                        try:
                            self.get(file, os.path.join(localdir,file), force)
                        except NoFileException, e:
                            print e
                        except error_perm, e:
                            print e
                        except FileExistsException, e:
                            print e
            finally:
                self.cwd(currentdir)

        def putdir(self, localdir, remotedir=None, force=False, recursive=False):
            """Upload local directory, recursivly if specified.
            
            Arguments:
            localdir -- local directory to upload
            remotedir -- remote location to store files
            force -- overwrite local files(True/False)
            recursive -- recursive upload(True/False)
            """
            
            if not self.connected: raise NoConnectionException, "No ftp connection present"

            if localdir[-1] == "/": localdir = localdir[0:-1]

            if remotedir == None: remotedir = localdir
            
            path = "/".join(remotedir.split("/")[0:-1])
            if path == "/": path = "./"
            remotedir = remotedir.split("/")[-1]

            if self.verbose: print "uploading %s/%s to %s" % (path, localdir, remotedir)

            currentdir = self.pwd()

            try:
                self.cwd(path)

                if remotedir in self.dirlist:
                    if not force:
                        raise FileExistsException, "remotedir %s allready exsists" % remotedir
                    else:
                        self.cwd(remotedir)
                else:
                    try:
                        self.sftp.mkdir(remotedir)
                        self.cwd(remotedir)
                    except Exception, e:
                        print e
                        return 1

                print self.sftp.getcwd()

                for file in os.listdir(localdir):
                    path = os.path.join(localdir,file)
                    if os.path.isfile(path):
                        if self.verbose: print "uploading %s" % path
                        self.put(path, file, force)
                    elif recursive and os.path.isdir(path):    
                        if self.verbose: print "uploading dir %s" % path
                        self.putdir(path, file, force, recursive)
            finally:
                self.cwd(currentdir)

        def rmdir(self, remotedir):
            """Recursively delte remote directory

            Arguments:
            remotedir -- remote directory to remove
            """
            
            if not self.connected: raise NoConnectionException, "No ftp connection present"

            if self.verbose: print "deleting %s" % remotedir

            if remotedir[-1] == "/": remotedir = remotedir[0:-1]

            path = "/".join(remotedir.split("/")[0:-1])
            if path == "/": path = "./"
            remotedir = remotedir.split("/")[-1]

            if remotedir == "": (remotedir,path) = (path,remotedir)

            currentdir = self.pwd()

            try:
                self.cwd(os.path.join(path,remotedir))
                for file in self.dirlist:
                    if self.verbose: print "deleting file %s" % file
                    if str(self.sftp.stat(file))[0] == '-':
                        self.delete(file) 
                    else:
                        self.rmdir(file)
                self.cwd(os.path.join(self.sftp.getcwd(), ".."))
                self.sftp.rmdir(remotedir)
            finally:
                self.cwd(currentdir)
                
            self.__setDirlist()

        def delete(self, remotefile):
            """Delete remotefile.

            Arguments:
            remotefile -- file on the server to remove
            """ 
            
            (remotedir, file) = os.path.split(remotefile)
            if remotedir == "": remotedir = "."
            
            currentdir = self.pwd()

            self.cwd(remotedir)

            if file in self.dirlist:
                if self.verbose: print "Deleting %s" % remotefile
                try:
                    self.sftp.remove(file)
                except IOError, e:
                    if self.verbose: print "IOError while deleting %s" % remotefile
                    raise NoFileException, "Error removing %s on server" % remotefile
            else:
                self.cwd(currentdir)
                raise NoFileException, "No file %s on server" % remotefile

            self.cwd(currentdir)

        def size(self, remotefile):
            """Get filesize of remotefile.

            Arguments:
            remotefile -- file on the server to check size of
            """
            if not self.connected: raise NoConnectionException, "No ftp connection present"

            (remotedir, remote) = os.path.split(remotefile)

            currentdir = self.pwd()

            try:
                self.cwd(remotedir)
                for file in self.dirlist:
                    if remote == file:
                        if self.verbose: print "found file %s, size: %s" % (remote,self.sftp.stat(file).st_size)
                        return self.sftp.stat(file).st_size
            finally:
                self.cwd(currentdir)

            if self.verbose: print "File not found: %s" % remotefile

            raise NoFileException, "File not found, %s" % remotefile

        def __setDirlist(self):
            """Set directory listing of current directory"""
            self.dirlist = self.sftp.listdir()
            
            self.dirlist = [a for a in self.dirlist if a[0]!='.']

        def time(self, remotefile):
            """Get creation time of remotefile in seconds since Epoch.
            
            Arguments:
            remotefile -- remotefile to get creation time of
            """ 
            
            monthdict = {'Jan': 1,'Feb': 2,'Mar': 3,'Apr': 4,'May': 5,'Jun': 6,'Jul': 7,'Aug': 8,'Sep': 9,'Oct': 10,'Nov': 11,'Dec':12} 

            if not self.connected: raise NoConnectionException, "No ftp connection present"
            (remotedir, remote) = os.path.split(remotefile)

            currentdir = self.pwd()

            time = None
            try:
                self.cwd(remotedir)
                for file in self.dirlist:
                    if remote == file:
                        time = self.sftp.stat(file).st_mtime
            finally:
                self.cwd(currentdir)

            return time
